%entry point for the AoT based 3D car template learning project.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Section 0, configuration and parameters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
configPath = './config_ver6' % path for root node configuration file

source_img_path_list = cell(0);
source_img_path_list{1} = '../UCLACar/Car_Frontal Side View';
source_img_path_list{2} = '../UCLACar/Car_Frontal View';
source_img_path_list{3} = '../UCLACar/Car_Rear Side View';
source_img_path_list{4} = '../UCLACar/Car_Side View';
source_img_path_list{5} = '../UCLAHighTilt';

resPath = '../results'

cachePath = '../cache/0'
if ~exist(cachePath)
    mkdir(cachePath)
end

resTrees=cell(0);% corresponds to  list of volOR node
file_name = [resPath '/template.con*'];
if exist(file_name)
    delete(file_name);
end


addpath('./libFilter');

makeLibFilter;
mex  outputTanAppearance.cpp projector.cpp shapeAndResponse.cc svgDrawer.cpp
mex  computeTanScore.cpp projector.cpp shapeAndResponse.cc svgDrawer.cpp
mex  computeColorModel.cpp colorAndResponse.cpp projector.cpp

default_filter_param;

loadAndFilterImages;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Section 1, set up the AoT for learning
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

generateTanConfiguration(configPath);
config_files = dir([configPath '/*.tan']);
% collection of AoT and parameters. Each tree here corresponds to a part in model
AoTs = [];
AoTParams = [];
for i_file = 1:length(config_files)
    AoT_param = parseConfigFile(configPath,config_files(i_file).name);
    AoT_param.name = config_files(i_file).name(4:end-4);
    AoTParams=[AoTParams;( AoT_param)];
end


% color_gain_th will be meaning less, as we will select all of them
color_gain_th = zeros(12,1)+0.045;
% part names are 
% 1-4: front_window head left_door lid
% 5-8: rear_window right_door roof tail
% 9-12 four wheels

tree_index =[1:12];
vol_and_number = [1,2,1,1,1,1,1,2,2,2,2,2];

for i_tree_ind = length(tree_index):-1:1
    i_tree = tree_index(i_tree_ind);
    tic
    corner = AoTParams(i_tree).corner;
    dim = AoTParams(i_tree).size;
  
    AoT = createTree(corner,dim(1),dim(2),dim(3),...
        AoTParams(i_tree).asp_low_limit,AoTParams(i_tree).asp_high_limit,...
        AoTParams(i_tree).rotationMatrix, AoTParams(i_tree).swt,AoTParams(i_tree).stp,...
	AoTParams(i_tree).name);
   
    AoT.name = AoTParams(i_tree).name;
    toc
    disp([ 'computing appearance term for : ' AoT.name]);
    AoTs = [AoTs;AoT];
    file_name = [ 'AoTree_' num2str(i_tree,'%02d') '.mat'];
    save(file_name,'AoT');
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % Section 2, compute the response for each tan in terminal node
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    tic
	[score_cell arg_max_cell]=computeTanScore(n_image,nOrient,nLength,h,Sx,Sy,Zx,Zy,M2_maps,PMat,viewMat,AoT);
	
    for iLayer = 1:numel(AoT.termLayer)
        %disp(['iLayer: ' num2str(iLayer) ' of iTree: ' num2str(i_tree)])
        num_nodes = numel(AoT.termLayer{iLayer});
        if num_nodes ==0
            continue;
        end
		for iNode = 1:numel(AoT.termLayer{iLayer})
			AoT.termLayer{iLayer}(iNode).score_map = score_cell{iLayer}{iNode};
			AoT.termLayer{iLayer}(iNode).arg_max_map = arg_max_cell{iLayer}{iNode};
		end
    end
    toc
	clear score_cell
	clear arg_max_cell
	disp('saving AoT with Gain');  
%    file_name= ['AoTWithGain_' num2str(i_tree) '.mat'];
%    save(file_name,'AoT');
%	AoT = [];

%    file_name= ['AoTWithGain_' num2str(i_tree) '.mat'];
%    load(file_name,'AoT');
    n_vol_and = vol_and_number(i_tree);
    AoT.vol_and_cluster=ceil(n_vol_and*rand(n_image,1));
    
    for iEM = 5:5;
        disp(['iteration ' num2str(iEM) ', tree ' num2str(i_tree)])
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Section 3, compute the IG for each tan in each cluster
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
     
        disp('compute IG for each tan');
        for iLayer = 1:numel(AoT.termLayer)
            for iNode = 1:numel(AoT.termLayer{iLayer})
                [IG, ind]=computeTanIG(AoT.termLayer{iLayer}(iNode).score_map,1,...
                    AoT.vol_and_cluster, n_vol_and);
                AoT.termLayer{iLayer}(iNode).score = IG;
                AoT.termLayer{iLayer}(iNode).shape_ind = int32(ind);
            end
        end
  
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Section 4, assume IG is computed, use DP to select node
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%     
        
        disp('extracting solution from tree ');
        resTree = extractSolution(n_vol_and,AoT.andLayer,AoT.orLayer,AoT.termLayer);
        out_path = [resPath '/tree_' num2str(i_tree,'%02d')];
	if ~exist(out_path,'dir')
		mkdir(out_path);
	end
	file_name = [ resPath '/tree_' num2str(i_tree,'%02d') '/solTree_' num2str(i_tree,'%02d') '_itr' num2str(iEM,'%02d') '.mat'];
        save(file_name,'resTree');
        
        resTrees{i_tree}=resTree';
        
	if iEM==5
        disp('visualize the trees')
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Visualize the trees
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        
        %disp('loading S1 and S2 maps...');
        %load([cachePath '/S1maps.mat']);
        %load([cachePath '/S2maps.mat']);
        load([cachePath '/filters.mat']);
        load([cachePath '/color_features.mat']);
        n_vol_and = numel(resTree);
	for i_vol_and = 1:n_vol_and
            current_res_tree = resTree{i_vol_and};
            for iLayer = 1:numel(current_res_tree.termLayer)
                for iNode = 1:numel(current_res_tree.termLayer{iLayer})
                    if ~strcmp(current_res_tree.termLayer{iLayer}(iNode).type,'Term')
                        continue;
                    end
                    out_path = [resPath '/tree_' num2str(i_tree,'%02d') '/itrEM' num2str(iEM)  '/cluster_'  num2str(i_vol_and) '/tan_' num2str(iLayer) '_' num2str(iNode)];
                    if~exist(out_path,'dir')
                        mkdir(out_path)
                    end
                    plane_name = ['Tree_' num2str(i_tree,'%02d') '/tan_' num2str(iLayer) '_' num2str(iNode)];
                    % output shape file
                    if_plot = double(AoT.vol_and_cluster==i_vol_and);
                    term_node=current_res_tree.termLayer{iLayer}(iNode);
		    outputTanAppearance(n_image,nOrient,nLength,h,Sx,Sy,Zx,Zy,M2_maps,...
                        AoT,term_node,...
                        PMat,viewMat,numel(out_path),out_path,if_plot);

					% remove redundant shape file
					tmp_file_list = dir([out_path '/*.pshape']);
					tmp_i_type = sscanf(tmp_file_list(1).name,'image_%d_shape_%d_deform%d.pshape');
					tmp_i_type = tmp_i_type(2);
					tmp_file_list =dir([out_path '/*.shape'])
					for tmp_i_file = 1:length(tmp_file_list)
						if ~strcmp(tmp_file_list(tmp_i_file).name,['canonical_' num2str(tmp_i_type) '.shape']);
							delete([out_path '/' tmp_file_list(tmp_i_file).name]);
						else
							copyfile([out_path '/' tmp_file_list(tmp_i_file).name],[out_path '/learned_shape.txt']);
                        end
					end
                    % output color model
					disp('compute color model...')
                    arg_max_map = term_node.arg_max_map;
                    tmp_ind= sub2ind(size(arg_max_map),(1:n_image)',double(term_node.shape_ind));
                    deformation = arg_max_map(tmp_ind);
                    deformation((~if_plot))=-1e9;
                    file_name = [out_path '/color_model.color'];
					gain_th = color_gain_th(i_tree);
                    computeColorModel(n_image,nOrient,nLength,2*h+1,...
                        Sx,Sy,Zx,Zy,color_l_maps,color_a_maps,color_b_maps,...
                        PMat,viewMat,AoT,term_node,deformation,...
                       gain_th,file_name);
					% export R T swt of root volume to current folder
					file_name = [out_path '/' 'rootParam.mat'];
					R = single(AoT.rotationMat);
					T = single(AoT.origin);
					swt = AoT.swt;
					save(file_name,'R','T','swt');
                    % convert shape file to image file
                    file_list = dir([out_path '/*.shape']);
                    for i_file =1:length(file_list)
                        img = drawSVGtoImage([out_path '/' file_list(i_file).name]);
                        img = 255 -img;
                        [img_color alpha]= drawColorModel([out_path '/color_model.color'],img);
                        imwrite(img_color,[out_path '/' file_list(i_file).name '.png']);
                    end
                    if iEM==5
                    file_list = dir([out_path '/*.pshape']);
                    for i_file = 1:length(file_list)
                        img = drawSVGtoImage([out_path '/' file_list(i_file).name]);
                        imwrite(255-img,[out_path '/' file_list(i_file).name '.png']);
                        ind_img = sscanf(file_list(i_file).name,'image_%d_shape_%d_deform_%d_.pshape');
                        source_img = I_gray{ind_img(1)+1};
                        img = 0.5*(255-img) + 0.5*source_img;
                        imwrite(img,[out_path '/' file_list(i_file).name '.impose.png']);
                    end
                    
                    % export to sketch up about the plane and tan parameters
                    exportToSketchUp(AoT,current_res_tree.termLayer{iLayer}(iNode),...
                        current_res_tree.termLayer{iLayer}(iNode).plane,...
                        plane_name,out_path);
                    
                    % export this file to the template configuration
                    file_name = [resPath '/template.config'];
                    file = fopen(file_name,'a+');
                    fprintf(file,'%s\n',out_path);
                    fprintf(file,'%d\n',current_res_tree.termLayer{iLayer}(iNode).shape_ind(find(if_plot==true,1,'first')));
                    fclose(file);
                    
                    % back track further to S2 maps and S1 maps
                    drawArgM2Lines;
		    end
                end % iNode
            end %iLayer
        end % i_vol_and
	end % if display
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % Section: 5 step, to impute the new cluster index
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        disp('Estep');
        Estep;
    end% iEM
    file_name = [ 'solTree_' num2str(i_tree,'%02d') '_final.mat'];
    save(file_name,'resTree');
    AoT=[];
    clear AoT
 
end% i_tre
clear
